#coding:utf-8

BLACK = "black"
RED = "red"

class RbtNode(object):
    def __init__(self,val):
        self.val = val
        self.parent = None
        self.left = None
        self.right = None
        self.color = RED


    def setParent(self,node):
        self.parent = node

    def setColor(self,color):
        self.color = color

class RBT(object):
    def __init__(self):
        self.root = None
        self.find = 0
        self.nodes = []
        self.leafs = []
        self.blanknum = 0

    def insert(self,val):
         node = RbtNode(val)
         self.nodes.append(node)

         if(self.root is None):
             self.insert_fixup_2(node)
             return

         ite = self.root
         pnode = ite
         while ite is not None:
             pnode = ite
             if(ite.val > val):
                 ite = ite.left
             else:
                 ite = ite.right

         node.parent = pnode
         if(node.val > pnode.val):
             pnode.right = node
         else:
             pnode.left = node

         if(pnode.color is not BLACK):
             self.insert_fixup_2(node)
         else:
             return

    #  按照《算法导论》的方法进行调整设计
    def insert_fixup_2(self,node):
        if node.parent is None:
            self.root = node
            node.setColor(BLACK)
            return

        pnode = node.parent
        if(pnode.color is BLACK):
            return

        gnode = pnode.parent
        if gnode is None:
            return

        if pnode == gnode.left:
            uncle = gnode.right
            if uncle is not None and uncle.color == RED:
                uncle.color = BLACK
                pnode.color = BLACK
                gnode.color = RED
                self.insert_fixup_2(gnode)
            elif node == pnode.right:
                node = pnode
                self.left_rotate(node)
                node.parent.color = BLACK
                node.parent.parent.color = RED
                self.right_rotate(node.parent.parent)
            else:
                node.parent.color = BLACK
                node.parent.parent.color = RED
                self.right_rotate(node.parent.parent)
        else:
            uncle = gnode.left
            if uncle is not None and uncle.color == RED:
                uncle.color = BLACK
                pnode.color = BLACK
                gnode.color = RED
                self.insert_fixup_2(gnode)
            elif node == pnode.left:
                node = pnode
                self.right_rotate(node)
                node.parent.color = BLACK
                node.parent.parent.color = RED
                self.left_rotate(node.parent.parent)
            else:
                node.parent.color = BLACK
                node.parent.parent.color = RED
                self.left_rotate(node.parent.parent)

    #  错误的插入调整方法
    def insert_fixup(self,node):
        print("inser fix up")
        if node.parent is None:
            self.root = node
            node.setColor(BLACK)
            return

        pnode = node.parent
        if(pnode.color is BLACK):
            return

        gnode = pnode.parent
        if gnode is None:
            return
        leftc = False
        if(pnode.left == node):
            leftc = True

        unode = None
        if(gnode.left == pnode):
            unode = gnode.right
        else:
            unode = gnode.left

        if(unode is not  None and unode.color is RED):
            pnode.color = BLACK
            if unode is not None:
                unode.color = BLACK
            gnode.color = RED
            self.insert_fixup(gnode)
        elif(unode is None or unode.color is BLACK):
           cnode = None
           if(leftc):
               pnode.color = BLACK
               gnode.color = RED

               if(gnode.left == pnode):
                   cnode = gnode.left.right
                   self.right_rotate(gnode)
               else:
                   cnode = gnode.right.left
                   self.left_rotate(gnode)
               if cnode is not None:
                    self.insert_fixup(cnode)
           else:
               cnode = pnode.right.left
               self.left_rotate(pnode)

               if (gnode.left == pnode):
                   cnode = gnode.left.right
                   self.right_rotate(pnode)
               else:
                   cnode = gnode.right.left
                   self.left_rotate(gnode)


           if cnode is not None:
               self.insert_fixup(cnode)
        else:
            pass

    # left rotate
    def left_rotate(self,node):
        print("left rotate")
        rcnode = node.right
        if(rcnode is None):
            return

        rclnode = rcnode.left
        if(rclnode is not None):
            rclnode.parent = node
        node.right = rclnode

        pnode = node.parent
        rcnode.parent = pnode
        if(pnode.left == node):
            pnode.left = rcnode
        else:
            pnode.right = rcnode

        rcnode.left = node
        node.parent = rcnode


    def right_rotate(self,node):
        print("right rotate")
        lcnode = node.left
        if(lcnode is None):
            return

        lcrnode = lcnode.right
        if(lcrnode is not None):
            lcrnode.parent = node
        node.left = lcrnode

        pnode = node.parent
        if(pnode is None):
            lcnode.parent = None
            self.root = lcnode
        else:
            lcnode.parent = pnode
            if(pnode.left == node):
                pnode.left = lcnode
            else:
                pnode.right = lcnode
        lcnode.right = node
        node.parent= lcnode

    def find_node(self,val):
        curNode = self.root
        while curNode is not None:
            self.find += 1
            print "current node val = ", curNode.val
            if val == curNode.val:
                print "yes ,find", val, "findnum=", self.find
                return
            elif val < curNode.val:
                curNode = curNode.left
            else:
                curNode = curNode.right
        print "not find,sorry", "findnum=", self.find
        return

    def check_rbt(self):
        for i in xrange(len(self.nodes)):
            self.check_leaf(self.nodes[i])

        self.check_color()

        for node in self.leafs:
            print("leaf node =%d"%node.val)
            self.get_blank_num(node)

    def check_leaf(self,node):
        if node.left is None and node.right is None:
            self.leafs.append(node)

    def check_color(self):
        print("check color begin...")
        for node in self.nodes:
            if node.left is not None and node.color is RED:
                leftn = node.left
                if leftn.color is RED:
                    print("red node have red left child....,node.val =%d"%(node.val))

            if node.right is not None and node.color is RED:
                rightn = node.right
                if rightn.color is RED:
                    print("red node have red right child....,node.val =%d"%(node.val))
        print("check color end....")

    def get_blank_num(self,node):
        curNode = self.root
        num = 0
        while curNode is not None:
            print ("curNode.color=%s"%curNode.color)
            if curNode.color is BLACK:
                num += 1

            if node == curNode:
                if(self.blanknum == 0):
                    self.blanknum = num
                else:
                    if(self.blanknum != num):
                        print("black num is not equal...,num=%d,self.blanknum=%d,node.val=%d"%(num,self.blanknum,node.val))
                return
            elif node.val < curNode.val:
                curNode = curNode.left
            else:
                curNode = curNode.right
        return

if __name__ == "__main__":
    rbt = RBT()
    b = []
    # TArr = [127,281,431,11,15,232,292,461,314,445]
    # TArr = [80,60,120,40,70,140,20,50]
    TArr = [80, 60, 120, 40, 70, 140, 20, 50, 100, 52, 96, 87, 25, 41, 96, 25, 14, 23, 21, 47, 578, 14, 9999, 621, 554]
    # TArr = [80, 60, 120, 40, 70, 140, 20, 50, 100, 52, 96, 87, 25, 41, 96, 25, 14]

    for i in range(len(TArr)):
        m = TArr[i]
        print("i insert val=%d" % m)
        rbt.insert(m)
        b.append(m)
    rbt.insert(35)
    TArr = TArr + [35]
    for i in range(len(TArr)):
        m = TArr[i]
        print("i begin val=%d"%m)
        rbt.find_node(m)
        print("i end")
    rbt.check_rbt()